package crypto.rsa

import apps.games.GameExecutionException
import crypto.random.secureRandom
import entity.User
import org.bouncycastle.crypto.AsymmetricBlockCipher
import org.bouncycastle.crypto.encodings.PKCS1Encoding
import org.bouncycastle.crypto.engines.RSAEngine
import org.bouncycastle.crypto.params.RSAKeyParameters
import org.bouncycastle.crypto.util.PrivateKeyFactory
import org.bouncycastle.crypto.util.PublicKeyFactory
import org.bouncycastle.jcajce.provider.asymmetric.rsa.BCRSAPublicKey
import org.bouncycastle.jce.ECNamedCurveTable
import org.bouncycastle.jce.provider.BouncyCastleProvider
import sun.misc.BASE64Decoder
import sun.misc.BASE64Encoder
import java.math.BigInteger
import java.security.KeyPair
import java.security.KeyPairGenerator
import java.security.Security
import javax.xml.bind.DatatypeConverter

/**
 * Created by user on 7/28/16.
 */
val ECParams = ECNamedCurveTable.getParameterSpec("secp256k1")

class RSAKeyManager(val KEY_LENGTH: Int = 1024) {
    private val userEncodeEngines = mutableMapOf<User, AsymmetricBlockCipher>()
    private val userPublicKeys = mutableMapOf<User, String>()
    private val userDecodeEngines = mutableMapOf<User, AsymmetricBlockCipher>()
    private val userPrivateKeys = mutableMapOf<User, String>()
    val engine = PKCS1Encoding(RSAEngine())
    private lateinit var keyPair: KeyPair

    init {
        Security.addProvider(BouncyCastleProvider())
        reset()
    }

    /**
     * Get generated public key(string representation_
     */
    fun getPublicKey(): String {
        val b64encoder = BASE64Encoder()
        return b64encoder.encode(keyPair.public.encoded)
    }

    /**
     * get public exponent
     */
    fun getExponent(): BigInteger {
        return (keyPair.public as BCRSAPublicKey).publicExponent
    }

    /**
     * get public modulus value
     */
    fun getModulus(): BigInteger {
        return (keyPair.public as BCRSAPublicKey).modulus
    }

    /**
     * Get private key
     *
     * DISCLAIMER: revealing private key
     * might me a huge vulnerability
     */
    fun getPrivateKey(): String {
        val b64encoder = BASE64Encoder()
        return b64encoder.encode(keyPair.private.encoded)
    }

    /**
     * register public key generated by [getPublicKey]
     * for given user
     * @param user - User to be added
     * @param publicKey - public key of given user
     */
    fun registerUserPublicKey(user: User, publicKey: String) {
        if (user in userPublicKeys) {
            if (userPublicKeys[user] != publicKey) {
                throw GameExecutionException("Changing keys for user is not allowed: ${userPublicKeys[user]} \n $publicKey")
            }
            return
        }
        val b64decoder = BASE64Decoder()
        val keyParam = PublicKeyFactory.createKey(b64decoder.decodeBuffer(publicKey))
        val cypher: AsymmetricBlockCipher = PKCS1Encoding(RSAEngine())
        cypher.init(true, keyParam)
        userPublicKeys[user] = publicKey
        userEncodeEngines[user] = cypher
    }

    /**
     * register private key generated by [getPrivateKey]
     * for given user
     * @param user - User to be added
     * @param privateKey - private key of given user
     */
    fun registerUserPrivateKey(user: User, privateKey: String) {
        if (user in userPrivateKeys) {
            if (userPrivateKeys[user] != privateKey) {
                throw GameExecutionException("Changing keys for user is not allowed")
            }
            return
        }
        val b64decoder = BASE64Decoder()
        val keyParam = PrivateKeyFactory.createKey(b64decoder.decodeBuffer(privateKey))
        val cypher: AsymmetricBlockCipher = PKCS1Encoding(RSAEngine())
        cypher.init(false, keyParam)
        userDecodeEngines[user] = cypher
    }

    /**
     * Encode message with given public rsa parameters
     */
    fun encodeWithParams(mod: BigInteger,
                         exp: BigInteger,
                         msg: String): String {
        val cypher: AsymmetricBlockCipher = PKCS1Encoding(RSAEngine())
        cypher.init(true, RSAKeyParameters(false, mod, exp))
        return encodeWithEngine(cypher, msg)
    }

    /**
     * Encode message with given public rsa parameters
     */
    fun encodeWithParams(mod: BigInteger,
                         exp: BigInteger,
                         msg: ByteArray): String {
        val cypher: AsymmetricBlockCipher = PKCS1Encoding(RSAEngine())
        cypher.init(true, RSAKeyParameters(false, mod, exp))
        return encodeWithEngine(cypher, msg)
    }


    /**
     * encode message with users public key
     */
    fun encodeForUser(user: User, msg: String): String {
        val e = userEncodeEngines[user] ?: throw NoSuchUserException("No engine for user ${user.name}")
        return encodeWithEngine(e, msg)
    }

    /**
     * encode message with given rsa engine
     */
    private fun encodeWithEngine(engine: AsymmetricBlockCipher,
                                 msg: String): String {
        return encodeWithEngine(engine, msg.toByteArray())
    }

    /**
     * encode message with given rsa engine
     */
    private fun encodeWithEngine(engine: AsymmetricBlockCipher,
                                 bytes: ByteArray): String {
        var len = engine.inputBlockSize
        val res = StringBuilder()
        for (i in 0..bytes.size - 1 step len) {
            len = Math.min(len, bytes.size - i)
            res.append(toHexString(engine.processBlock(bytes, i, len)))
        }
        return res.toString()
    }


    /**
     * try to decode message with our own key
     * @param msg - message to decode
     * @throws InvalidCipherTextException  - if message has incorrect format
     */
    fun decodeString(msg: String): String {
        return decodeStringWithEngine(engine, msg)
    }

    /**
     * decode message of bytes
     */
    fun decodeBytes(msg: String): ByteArray {
        return decodeBytesWithEngine(engine, msg)
    }

    /**
     * try to decode message for user whose private key we know
     * @param msg - message to decode
     * @throws InvalidCipherTextException  - if message has incorrect format
     */
    fun decodeForUser(user: User, msg: String): String {
        val engine = userDecodeEngines[user] ?: throw IllegalArgumentException("No private key known for ser ${user.name}")
        return decodeStringWithEngine(engine, msg)
    }


    /**
     * Given string message try to decode it with given block cypher
     *
     * @param msg - string representing encoded message
     * @param engine - decoding cypher
     *
     * @return - decoded string
     */
    private fun decodeStringWithEngine(engine: AsymmetricBlockCipher,
                                       msg: String): String {
        var len = engine.inputBlockSize
        val res = StringBuilder()
        val bytes = toByteArray(msg)
        for (i in 0..bytes.size - 1 step len) {
            len = Math.min(len, bytes.size - i)
            res.append(String(engine.processBlock(bytes, i, len)))
        }
        return res.toString()
    }

    /**
     * Given bytes of message try to decode it with given block cypher
     *
     * @param msg - string representing encoded message
     * @param engine - decoding cypher
     *
     * @return - decoded string
     */
    private fun decodeBytesWithEngine(engine: AsymmetricBlockCipher,
                                      msg: String): ByteArray {
        var len = engine.inputBlockSize
        val res = mutableListOf<Byte>()
        val bytes = toByteArray(msg)
        for (i in 0..bytes.size - 1 step len) {
            len = Math.min(len, bytes.size - i)
            res.addAll(engine.processBlock(bytes, i, len).toList())
        }
        return res.toByteArray()
    }

    private fun toHexString(array: ByteArray): String {
        return DatatypeConverter.printHexBinary(array)
    }

    private fun toByteArray(s: String): ByteArray {
        return DatatypeConverter.parseHexBinary(s)
    }


    /**
     * reset private/public keys as well as list of known enginies
     */
    fun reset() {
        val keyGen = KeyPairGenerator.getInstance("rsa", "BC")
        keyGen.initialize(KEY_LENGTH, secureRandom)
        keyPair = keyGen.genKeyPair()
        engine.init(false, PrivateKeyFactory.createKey(keyPair.private.encoded))
        userEncodeEngines.clear()
        userDecodeEngines.clear()
        userPrivateKeys.clear()
        userPublicKeys.clear()
    }

}